{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pip install llama-index-llms-reka\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To obtain API key, please visit [https://platform.reka.ai/](https://platform.reka.ai/)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Chat completion\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " The capital city of France is Paris. It is located in the northern part of the country along the Seine River and is known for its rich history, iconic landmarks such as the Eiffel Tower, Louvre Museum, and Notre-Dame Cathedral, as well as its vibrant culture and cuisine.\n",
      "\n",
      "\n",
      " The capital city of France is Paris. Located in the northern part of the country along the Seine River, Paris is known for its rich history, iconic landmarks such as the Eiffel Tower, Louvre Museum, and Notre-Dame Cathedral, as well as its vibrant culture, fashion, and cuisine. It serves as the political, economic, and cultural center of France.\n",
      "\n",
      "\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "from llama_index.llms.reka import RekaLLM\n",
    "from llama_index.core.base.llms.types import ChatMessage, MessageRole\n",
    "\n",
    "# Initialize the Reka LLM\n",
    "api_key = os.getenv(\"REKA_API_KEY\")\n",
    "reka_llm = RekaLLM(\n",
    "    model=\"reka-flash\",\n",
    "    api_key=api_key,\n",
    ")\n",
    "\n",
    "# Chat completion\n",
    "messages = [\n",
    "    ChatMessage(role=MessageRole.SYSTEM, content=\"You are a helpful assistant.\"),\n",
    "    ChatMessage(role=MessageRole.USER, content=\"What is the capital of France?\"),\n",
    "]\n",
    "response = reka_llm.chat(messages)\n",
    "print(response.message.content)\n",
    "\n",
    "# Text completion\n",
    "prompt = \"The capital of France is\"\n",
    "response = reka_llm.complete(prompt)\n",
    "print(response.text)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Streaming example\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " Here are the first 5 planets in our solar system, listed in order from the closest to the Sun to the farthest:\n",
      "\n",
      "1. Mercury\n",
      "2. Venus\n",
      "3. Earth\n",
      "4. Mars\n",
      "5. Jupiter\n",
      "\n",
      " Here are the first 5 planets in our solar system, listed in order from the closest to the Sun to the farthest:\n",
      "\n",
      "1. Mercury\n",
      "2. Venus\n",
      "3. Earth\n",
      "4. Mars\n",
      "5. Jupiter\n",
      "\n"
     ]
    }
   ],
   "source": [
    "# Streaming chat completion\n",
    "messages = [\n",
    "    ChatMessage(role=MessageRole.SYSTEM, content=\"You are a helpful assistant.\"),\n",
    "    ChatMessage(\n",
    "        role=MessageRole.USER, content=\"List the first 5 planets in the solar system.\"\n",
    "    ),\n",
    "]\n",
    "for chunk in reka_llm.stream_chat(messages):\n",
    "    print(chunk.delta, end=\"\", flush=True)\n",
    "\n",
    "# Streaming text completion\n",
    "prompt = \"List the first 5 planets in the solar system:\"\n",
    "for chunk in reka_llm.stream_complete(prompt):\n",
    "    print(chunk.delta, end=\"\", flush=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Async use cases (chat/completion)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " The largest planet in our solar system is Jupiter. It is a gas giant with a diameter of about 139,822 kilometers (86,881 miles) at its equator, making it more than twice the size of Earth. Jupiter's massive size accounts for over 300 times the mass of all other planets in the solar system combined.\n",
      "\n",
      "\n",
      " The largest planet in our solar system is Jupiter. It is a gas giant with a diameter of about 86,881 miles (139,822 kilometers) at its equator, making it more than twice the size of Earth. Jupiter's mass is also more than 300 times that of Earth, and it contains over 70% of the total mass of all the planets in the solar system combined. Its prominent features include its Great Red Spot, a massive storm that has been raging for at least 400 years, and its extensive system of rings and more than 80 known moons.\n",
      "\n",
      "\n",
      " The first 5 elements in the periodic table, in order of increasing atomic number, are:\n",
      "\n",
      "1. Hydrogen (H) - Symbol: H, Atomic Number: 1\n",
      "2. Helium (He) - Symbol: He, Atomic Number: 2\n",
      "3. Lithium (Li) - Symbol: Li, Atomic Number: 3\n",
      "4. Beryllium (Be) - Symbol: Be, Atomic Number: 4\n",
      "5. Boron (B) - Symbol: B, Atomic Number: 5\n",
      "\n",
      " Here are the first 5 elements in the periodic table, listed in order of increasing atomic number:\n",
      "\n",
      "1. Hydrogen (H) - Atomic number 1\n",
      "2. Helium (He) - Atomic number 2\n",
      "3. Lithium (Li) - Atomic number 3\n",
      "4. Beryllium (Be) - Atomic number 4\n",
      "5. Boron (B) - Atomic number 5\n",
      "\n"
     ]
    }
   ],
   "source": [
    "async def main():\n",
    "    # Async chat completion\n",
    "    messages = [\n",
    "        ChatMessage(role=MessageRole.SYSTEM, content=\"You are a helpful assistant.\"),\n",
    "        ChatMessage(\n",
    "            role=MessageRole.USER,\n",
    "            content=\"What is the largest planet in our solar system?\",\n",
    "        ),\n",
    "    ]\n",
    "    response = await reka_llm.achat(messages)\n",
    "    print(response.message.content)\n",
    "\n",
    "    # Async text completion\n",
    "    prompt = \"The largest planet in our solar system is\"\n",
    "    response = await reka_llm.acomplete(prompt)\n",
    "    print(response.text)\n",
    "\n",
    "    # Async streaming chat completion\n",
    "    messages = [\n",
    "        ChatMessage(role=MessageRole.SYSTEM, content=\"You are a helpful assistant.\"),\n",
    "        ChatMessage(\n",
    "            role=MessageRole.USER,\n",
    "            content=\"Name the first 5 elements in the periodic table.\",\n",
    "        ),\n",
    "    ]\n",
    "    async for chunk in await reka_llm.astream_chat(messages):\n",
    "        print(chunk.delta, end=\"\", flush=True)\n",
    "\n",
    "    # Async streaming text completion\n",
    "    prompt = \"List the first 5 elements in the periodic table:\"\n",
    "    async for chunk in await reka_llm.astream_complete(prompt):\n",
    "        print(chunk.delta, end=\"\", flush=True)\n",
    "\n",
    "\n",
    "await main()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "llamaindex-reka",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
